import pandas as pd
from sklearn.preprocessing import LabelEncoder

def labelEncoder_handleCategoricalVariables():
    """
    使用LabelEncoder对分类变量进行编码
    :return:
    """
    # 读取数据
    train_df = pd.read_csv('../../data/train.csv')
    test_df = pd.read_csv('../../data/test2.csv')
    print(test_df.columns)
    print(test_df.shape)

    # 初始化LabelEncoder
    label_encoder = LabelEncoder()

    # 获取所有分类列（假设这些是分类变量）
    categorical_cols = ['Age','BusinessTravel','Department','DistanceFromHome','Education',
                        'EducationField','EnvironmentSatisfaction','Gender','JobInvolvement',
                        'JobLevel','JobRole','JobSatisfaction','MaritalStatus','MonthlyIncome',
                        'NumCompaniesWorked','OverTime','PercentSalaryHike','PerformanceRating',
                        'RelationshipSatisfaction','StockOptionLevel','TotalWorkingYears',
                        'TrainingTimesLastYear','WorkLifeBalance','YearsAtCompany','YearsInCurrentRole',
                        'YearsSinceLastPromotion','YearsWithCurrManager']
    # 先合并训练集和测试集，确保所有类别都被编码
    combined_df = pd.concat([train_df[categorical_cols], test_df[categorical_cols]])

    # # 对训练集和测试集的分类列进行编码
    for col in categorical_cols:
        # 在训练集上拟合并转换
        label_encoder.fit(combined_df[col])
        train_df[col] = label_encoder.transform(train_df[col])
        # 在测试集上转换（使用训练集的编码）
        test_df[col] = label_encoder.transform(test_df[col])

    # 保存处理后的数据
    train_df.to_csv('../../data/train_encoded.csv', index=False)
    test_df.to_csv('../../data/test_encoded.csv', index=False)

    print("数据处理完成，已保存为train_encoded.csv和test_encoded.csv")

if __name__ == '__main__':
    labelEncoder_handleCategoricalVariables()


